import os
import pickle
import numpy as np
import pandas as pd
import statsmodels.api as sm
from sklearn.preprocessing import MinMaxScaler,StandardScaler
import warnings
warnings.simplefilter("ignore")
import json

# 返回股票数据字典
def load_raw_data(data_list):
    all_dfs = {}
    for name in data_list:
        df=pd.read_csv(fr'D:\多任务选股_因子计算\factor_calculation\data_processed\data_{name}_processed.csv',header=0,index_col=0)
        df.index=pd.to_datetime(df.index,format='%Y-%m-%d')
        all_dfs[name]=df
    return all_dfs


#特定日期的股票数据字典
def load_raw_data1(data_list, start_date=None, end_date=None):
    all_dfs = {}
    for name in data_list:
        df = pd.read_csv(fr'D:\多任务选股_因子计算\factor_calculation\data_processed\data_{name}_processed.csv', header=0, index_col=0)
        df.index = pd.to_datetime(df.index, format='%Y-%m-%d')

        # 筛选指定日期范围的数据
        if start_date and end_date:
            df = df.loc[start_date:end_date]
        elif start_date:
            df = df.loc[start_date:]
        elif end_date:
            df = df.loc[:end_date]

        all_dfs[name] = df
    return all_dfs

# 保存因子数据为pkl
def save_factor(factor_df, factor_name,n):
    output_path=fr"D:\多任务选股_因子计算\factor_calculation\factor\{factor_name}.pkl"

    with open(output_path, 'wb') as f:
        pickle.dump(factor_df, f)

# 去极值
def winsorize(x:pd.DataFrame,method='med', qrange=(0.03, 0.97), scale=5.0):
    if method == 'percentile':
        max_value = x.quantile(q=qrange[1], axis=1)
        min_value = x.quantile(q=qrange[0], axis=1)
    elif method == 'sigma':
        mean_value = x.mean(axis=1)
        n_std_value = x.std(axis=1) * scale
        max_value = mean_value.add(n_std_value, axis=0)
        min_value = mean_value.sub(n_std_value, axis=0)
    elif method == 'med':
        median_value = x.median(axis=1)
        n_mad = np.absolute(x.sub(median_value, axis=0)).median(axis=1) * scale
        max_value = median_value.add(n_mad)
        min_value = median_value.sub(n_mad)
    return x.clip(lower=min_value, upper=max_value, axis=0)

# 对数市值中性化
def neutralize_market_log(x:pd.DataFrame, market_value:pd.DataFrame):
    factor_dropna = x.dropna(axis = 0, how = 'all')
    # display(factor_dropna)
    log_market=np.log(market_value).reindex(x.index).replace({-np.inf:np.nan,np.inf:np.nan})
    factor_neu = factor_dropna.apply(
        lambda x: sm.OLS(x, sm.add_constant(log_market.loc[x.name]), missing= 'drop').fit().resid, 
        axis = 1)
    return factor_neu

# 标准化
def standardlize(x:pd.DataFrame,method='z-score'):
    if method == 'max_min':
        return x.sub(x.min(axis=1), axis=0).div(x.max(axis=1).sub(x.min(axis=1)), axis=0)
    elif method == 'z-score':
        return x.sub(x.mean(axis=1), axis=0).div(x.std(axis=1), axis=0)

